import math
import torch
import torch.nn
import torch.optim
import torchvision
import numpy as np
import model
import config as c
import datasets
import modules.Unet_common as common

"""
模型训练完成后，进行测试
"""

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def load(name):
    """
       加载训练好的网络
    """
    state_dicts = torch.load(name) 
    network_state_dict = {k:v for k,v in state_dicts['net'].items() if 'tmp_var' not in k}
    net.load_state_dict(network_state_dict)
    try:
        optim.load_state_dict(state_dicts['opt'])
    except:
        print('Cannot load optimizer for some reason or other')


def gauss_noise(shape):
    """
        高斯噪音
    """
    noise = torch.zeros(shape).cuda()
    for i in range(noise.shape[0]):
        # 生成的随机值服从标准正态分布
        noise[i] = torch.randn(noise[i].shape).cuda()
    return noise


def computePSNR(origin,pred):
    """
        计算PSNR值
        origin:原始图片
        pred:预测后的图片
    """
    origin = np.array(origin)
    origin = origin.astype(np.float32)
    pred = np.array(pred)
    pred = pred.astype(np.float32)
    # 计算均方误差, 其实下面的除以1.0是可以去掉的
    mse = np.mean((origin/1.0 - pred/1.0) ** 2 )
    if mse < 1.0e-10:
      return 100
    return 10 * math.log10(255.0**2/mse) # PSNR


net = model.Model()
net.cuda()
# 测试模式，有必要初始化模型参数？
# model.init_model(net)
net = torch.nn.DataParallel(net, device_ids=c.device_ids)
# 筛选出需要计算梯度的参数
params_trainable = (list(filter(lambda p: p.requires_grad, net.parameters())))
# betas=c.betas：
# 设置Adam优化器中的两个超参数β1和β2。c.betas是一个包含两个值的元组，分别对应β1和β2。
# β1用于计算梯度的指数加权平均值，β2用于计算梯度平方的指数加权平均值。
# 这两个参数影响着优化器的动态调整学习率的能力。
optim = torch.optim.Adam(params_trainable, lr=c.lr, betas=c.betas, eps=1e-6, weight_decay=c.weight_decay)
weight_scheduler = torch.optim.lr_scheduler.StepLR(optim, c.weight_step, gamma=c.gamma)

load(c.MODEL_PATH + c.suffix)

# 开启评估模式
net.eval()

# 小波变换
dwt = common.DWT()
# 逆小波变换
iwt = common.IWT()


with torch.no_grad():
    for i, data in enumerate(datasets.get_testloader()):
        if i==1:
            break
        data = data.to(device) # 数据放到GPU上
        # 选择批次后一半的图片作为载体
        cover = data[data.shape[0] // 2:, :, :, :]
        # 选择批次前一半的图片作为秘密图片
        secret = data[:data.shape[0] // 2, :, :, :]
        # 用小波变换对图形做预处理
        cover_input = dwt(cover)
        secret_input = dwt(secret)
        input_img = torch.cat((cover_input, secret_input), 1)

        #################
        #    forward:   #
        #################
        output = net(input_img)
        # 将输出按通道分离
        # 隐写图像
        output_steg = output.narrow(1, 0, 4 * c.channels_in)
        # 图中的损失信息r
        output_z = output.narrow(1, 4 * c.channels_in, output.shape[1] - 4 * c.channels_in)
        # 逆小波变换
        steg_img = iwt(output_steg)
        backward_z = gauss_noise(output_z.shape)

        #################
        #   backward:   #
        #################
        output_rev = torch.cat((output_steg, backward_z), 1)
        bacward_img = net(output_rev, rev=True)
        secret_rev = bacward_img.narrow(1, 4 * c.channels_in, bacward_img.shape[1] - 4 * c.channels_in)
        # 恢复的秘密图像
        secret_rev = iwt(secret_rev)
        cover_rev = bacward_img.narrow(1, 0, 4 * c.channels_in)
        # 恢复的载体图像
        cover_rev = iwt(cover_rev)
        # 残差图像
        resi_cover = (steg_img - cover) * 20
        resi_secret = (secret_rev - secret) * 20

        # 保存图像
        torchvision.utils.save_image(cover, c.IMAGE_PATH_cover + '%.5d.png' % i)
        torchvision.utils.save_image(secret, c.IMAGE_PATH_secret + '%.5d.png' % i)
        torchvision.utils.save_image(steg_img, c.IMAGE_PATH_steg + '%.5d.png' % i)
        torchvision.utils.save_image(secret_rev, c.IMAGE_PATH_secret_rev + '%.5d.png' % i)




